PRINT_TRACE_TYPE0 = False
PRINT_TRACE_TYPE1 = False
PRINT_TRACE_TYPE2 = False
WIN_UTIL = 10000
WHITE = 'W'
BLACK = 'B'
ALLOW_SUBTREES = True
from math import ceil
from time import clock
from alpha_beta import INFINITY, AlphaBetaSearch
from LightState import LightQuoridorState

class AlphaBetaEnhancedSearch(AlphaBetaSearch):
    '''
    This search algorithm implements the limited-resource minimax tree search 
    algorithm with alpha-beta pruning for zero-sum games. This version cuts off 
    search by depth alone and uses an evaluation function (utility).
    this algorithm includes the enhancement of sorting the sons of each q_state in 
    such way the critical will be first.
    the search is limited with time_limit
    '''
    
    def __init__(self, player, global_max_depth, utility, time_limit, caching_on=False, w=1, succ_func=1):
        '''
        Constructor.
        '''
        self.end_time = 0
        self.time_limit = time_limit
        self.sorting_on = True
        self.caching_on = caching_on
        self.cache = {}
        self.cache_cnt = 0
        self.non_cache_cnt = 0
        self.w = w
        self.global_max_depth = global_max_depth
        self.look_for_final_line = False
        self.GO_BACK_ACTIONS = None
        if succ_func ==1 : self.SuccFunc = self.get_successors_by_utility
        else: self.SuccFunc = self.get_successors_by_relevance
        AlphaBetaSearch.__init__(self, player, global_max_depth, utility)
    
    def search(self, current_q_state):
        '''
        Search game to determine best action; use alpha-beta pruning.
        
        @param current_q_state: The current game q_state to start the search from.
        '''
        self.end_time = clock() + self.time_limit
        best_action = None
        if not self.look_for_final_line:
            next_rowsW = []
            next_rowsB = []
            for i in range(current_q_state.numPieces):
                if current_q_state.locationsW[i] is not None: next_rowsW.append(current_q_state.locationsW[i][0] +1)
                else: next_rowsW.append(None)
                if current_q_state.locationsB[i] is not None: next_rowsB.append(current_q_state.locationsB[i][0] -1)
                else: next_rowsB.append(None)
        else:
            next_rowsW = [current_q_state.brdSize-1]*current_q_state.numPieces
            next_rowsB = [0]*current_q_state.numPieces
        target_rowsW = next_rowsW
        target_rowsB = next_rowsB

        self.max_depth = 0
        best_value = -INFINITY
        while (self.max_depth < self.global_max_depth) and  (clock() < self.end_time) \
                     and (best_value < WIN_UTIL):            
            self.max_depth += 1
            self.num_nodes = 0
            best_value = -INFINITY
            cur_best_action = None
            if PRINT_TRACE_TYPE0: print "STARTED NEW TREE, ALLOWED DEPTH=%s" %  self.max_depth
            successors, utils = self.SuccFunc(current_q_state, 0, next_rowsW, next_rowsB, rev=True, values_only=False)
            
    #            if PRINT_TRACE_TYPE1:
    #                print "looking for max:"
    #                for action, q_state in successors:
    #                    print "action %s has hueristic %s" %(str(action).strip(), self.utility(q_state, target_rowsW, target_rowsB))
            timeout = False
            for i, (action, q_state) in enumerate(successors):
                if PRINT_TRACE_TYPE2: print "checking action %s" %(str(action).strip())
                value_fn = self._getValueFn(q_state)
                value = value_fn(q_state, best_value, INFINITY, 1, next_rowsW, next_rowsB, target_rowsW, target_rowsB, utils[i])
                if PRINT_TRACE_TYPE2: print "action %s can ensure utility %s" %(str(action).strip(), value)
                if value > best_value:
                    best_value = value
                    cur_best_action = action
                    best_rank = i+1
                if clock() >= self.end_time:
                    timeout = True
                    break
            if PRINT_TRACE_TYPE2: print "best rank=%s of %s" % (best_rank, len(successors))
            if not timeout or self.max_depth == 1:
                best_action = cur_best_action
#            print "Nodes in tree:", self.num_nodes
        if not self.look_for_final_line and self.GO_BACK_ACTIONS is not None and best_action in self.GO_BACK_ACTIONS:
            self.look_for_final_line = True
        return best_action
    

    def _maxValue(self, q_state, alpha, beta, depth, next_rowsW, next_rowsB, target_rowsW, target_rowsB, util):
        self.num_nodes += 1
#        if PRINT_TRACE_TYPE2 and depth > self.max_depth: 
#            print "sub-tree, _maxValue, depth=%s" % depth
#        print "minfunc, player turn:", q_state.curPlayer
        not_quiet = False
        if ALLOW_SUBTREES and depth == self.max_depth: 
            if self.player == WHITE and q_state.curPlayer == BLACK:
                for posB in q_state.locationsB:
                    if posB != None and posB[0] <= 1 and ((posB[1] == 0 and q_state.wallsH[posB[0]-1][posB[1]] == 0) or \
                        (posB[1] > 0 and (q_state.wallsH[posB[0]-1][posB[1]-1] == 0 and q_state.wallsH[posB[0]-1][posB[1]] == 0))):
                        not_quiet = True
            elif self.player == BLACK and q_state.curPlayer == WHITE:
                for posW in q_state.locationsW:
                    if posW != None and posW[0] >= q_state.brdSize - 2 and ((posW[1] == 0 and q_state.wallsH[posW[0]-1][posW[1]] == 0) or \
                        (posW[1] > 0 and (q_state.wallsH[posW[0]-1][posW[1]-1] == 0 and q_state.wallsH[posW[0]-1][posW[1]] == 0))):
                        not_quiet = True
        
        if (q_state.getWinner() is not None) or (clock() >= self.end_time) or (depth >= self.max_depth and not (not_quiet)):
            if util is None: util = self.utility(q_state, target_rowsW, target_rowsB, self.end_time)
#            util2 = self.utility(q_state, target_rowsW, target_rowsB, self.end_time)
#            print "util=%s, util2=%s" % (util , util2)
            if self.caching_on:
                subtree_depth = self.max_depth - depth
                if subtree_depth > 1 or True:
                    state_light = LightQuoridorState(q_state, target_rowsW, target_rowsB, len(q_state), alpha, beta)
                    if (subtree_depth in self.cache and state_light in self.cache[subtree_depth]):
                        #print 'FOUND IN CACHE!'
                        self.cache_cnt += 1
                        cache_val = self.cache[subtree_depth][state_light]
                        if cache_val != util:
                            print "max_func: CACHE NOT OK, depth=%s, values: %s, %s" % (subtree_depth, util, cache_val)
                            print state_light
                        #else: print "max_func: CACHE is OK, depth=%s, values: %s, %s" % (subtree_depth, my_value, cache_val)
                    if subtree_depth not in self.cache:
                        self.cache[subtree_depth] = {}
                    self.cache[subtree_depth][state_light] = util
            return util

        self.non_cache_cnt += 1

        value = -INFINITY
        
        successors, utils = self.SuccFunc(q_state, depth, next_rowsW, next_rowsB, rev=True, values_only=True)
#        print "len=", len(successors)
#            if PRINT_TRACE_TYPE1:
#                print "looking for max:"
#                for q_state in successors:
#                    print "hueristic=%s" %(self.utility(q_state, target_rowsW, target_rowsB))
        for i, successor in enumerate(successors):
            value_fn = self._getValueFn(successor)
            my_value = value_fn(successor, alpha, beta, depth + 1, next_rowsW, next_rowsB, target_rowsW, target_rowsB, utils[i])
            if self.caching_on:
                subtree_depth = self.max_depth - depth
                if subtree_depth > 1 or True:
                    state_light = LightQuoridorState(successor, target_rowsW, target_rowsB, len(successors), alpha, beta)
                    if (subtree_depth in self.cache and state_light in self.cache[subtree_depth]):
                        #print 'FOUND IN CACHE!'
                        self.cache_cnt += 1
                        cache_val = self.cache[subtree_depth][state_light]
                        if cache_val != my_value:
                            print "max_func: CACHE NOT OK, depth=%s, values: %s, %s" % (subtree_depth, my_value, cache_val)
                            print state_light
                        #else: print "max_func: CACHE is OK, depth=%s, values: %s, %s" % (subtree_depth, my_value, cache_val)
                    if subtree_depth not in self.cache:
                        self.cache[subtree_depth] = {}
                    self.cache[subtree_depth][state_light] = my_value
            value = max(value, my_value)
            if value >= beta:
                if PRINT_TRACE_TYPE1: print "max_func: cut after %s sons of %s at depth %d of %d" % (i+1, len(successors), depth, self.max_depth)
                return value
            alpha = max(alpha, value)
            if clock() >= self.end_time:
                break
           
        if PRINT_TRACE_TYPE1: print "max_func: cut after %s sons of %s at depth %d of %d" % (i+1, len(successors), depth, self.max_depth)
        return value
    
    def _minValue(self, q_state, alpha, beta, depth, next_rowsW, next_rowsB, target_rowsW, target_rowsB, util):
        self.num_nodes += 1
#        if PRINT_TRACE_TYPE2 and depth > self.max_depth: 
#            print "sub-tree, _minValue, depth=%s" % depth
#        print "minfunc, player turn:", q_state.curPlayer
        not_quiet = False 
        if ALLOW_SUBTREES and depth == self.max_depth: 
            if self.player == WHITE and q_state.curPlayer == BLACK:
                for posB in q_state.locationsB:
                    if posB != None and posB[0] <= 1 and ((posB[1] == 0 and q_state.wallsH[posB[0]-1][posB[1]] == 0) or \
                        (posB[1] > 0 and (q_state.wallsH[posB[0]-1][posB[1]-1] == 0 and q_state.wallsH[posB[0]-1][posB[1]] == 0))):
                        not_quiet = True
            elif self.player == BLACK and q_state.curPlayer == WHITE:
                for posW in q_state.locationsW:
                    if posW != None and posW[0] >= q_state.brdSize - 2 and ((posW[1] == 0 and q_state.wallsH[posW[0]-1][posW[1]] == 0) or \
                        (posW[1] > 0 and (q_state.wallsH[posW[0]-1][posW[1]-1] == 0 and q_state.wallsH[posW[0]-1][posW[1]] == 0))):
                        not_quiet = True
        
        if (q_state.getWinner() is not None) or (clock() >= self.end_time) or (depth >= self.max_depth and not (not_quiet)):
            if util is None: util = self.utility(q_state, target_rowsW, target_rowsB, self.end_time)
            if self.caching_on:
                subtree_depth = self.max_depth - depth
                if subtree_depth > 1 or True:
                    state_light = LightQuoridorState(q_state, target_rowsW, target_rowsB, len(q_state), alpha, beta)
                    if (subtree_depth in self.cache and state_light in self.cache[subtree_depth]):
                        #print 'FOUND IN CACHE!'
                        self.cache_cnt += 1
                        cache_val = self.cache[subtree_depth][state_light]
                        if cache_val != util:
                            print "max_func: CACHE NOT OK, depth=%s, values: %s, %s" % (subtree_depth, util, cache_val)
                            print state_light
                        #else: print "max_func: CACHE is OK, depth=%s, values: %s, %s" % (subtree_depth, my_value, cache_val)
                    if subtree_depth not in self.cache:
                        self.cache[subtree_depth] = {}
                    self.cache[subtree_depth][state_light] = util
#            print "util=", util, "blackMinDistance=", blackMinDistance
            return util

        self.non_cache_cnt += 1

        value = INFINITY
        
        successors, utils = self.SuccFunc(q_state, depth, next_rowsW, next_rowsB, rev=False, values_only=True)
        
#            if PRINT_TRACE_TYPE1:
#                print "looking for min:"
#                for q_state in successors:
#                    print "hueristic=%s" %(self.utility(q_state, target_rowsW, target_rowsB))
        for i, successor in enumerate(successors):
            value_fn = self._getValueFn(successor)
            my_value = value_fn(successor, alpha, beta, depth + 1, next_rowsW, next_rowsB, target_rowsW, target_rowsB, utils[i])
            if self.caching_on:
                subtree_depth = self.max_depth - depth
                if subtree_depth > 1 or True:
                    state_light = LightQuoridorState(successor, target_rowsW, target_rowsB, len(successors), alpha, beta)
                    if (subtree_depth in self.cache and state_light in self.cache[subtree_depth]):
                        #print 'FOUND IN CACHE!'
                        self.cache_cnt += 1
                        cache_val = self.cache[subtree_depth][state_light]
                        if cache_val != my_value:
                            print "min_func: CACHE NOT OK, depth=%s, values: %s, %s" % (subtree_depth, my_value, cache_val)
                            print state_light
                        #else: print "max_func: CACHE is OK, depth=%s, values: %s, %s" % (subtree_depth, my_value, cache_val)
                    if subtree_depth not in self.cache:
                        self.cache[subtree_depth] = {}
                    self.cache[subtree_depth][state_light] = my_value
            value = min(value, my_value)
            if value <= alpha:
                if PRINT_TRACE_TYPE1: print "min_func: cut after %s sons of %s at depth %d of %d" % (i+1, len(successors), depth, self.max_depth)
                return value
            beta = min(beta, value)
            if clock() >= self.end_time:
                break
       
        if PRINT_TRACE_TYPE1: print "min_func: cut after %s sons of %s at depth %d of %d" % (i+1, len(successors), depth, self.max_depth)
        return value

    def get_successors_by_utility(self, q_state, depth, next_rowsW, next_rowsB, rev=False, values_only=True):
        
        if not self.sorting_on:# or depth == self.max_depth - 1:
            successors = q_state.getSuccessors().items()
            utils = [None] * len(successors)
        else:
            if PRINT_TRACE_TYPE2: print "sorting"
#            successors = sorted(q_state.getSuccessors().items(), key=lambda suc: self.utility(suc[1], next_rowsW, next_rowsB, self.end_time), reverse=rev)
#            successors_2 = q_state.getSuccessors().items()
            successors = []
            notsorted = q_state.getSuccessors().items()
            for suc in notsorted:
                successors.append((self.utility(suc[1], next_rowsW, next_rowsB, self.end_time), suc))
                if clock() > self.end_time: break
            successors = sorted(successors, reverse=rev)
#            successors = sorted([(self.utility(suc[1], next_rowsW, next_rowsB, self.end_time), suc) for suc in q_state.getSuccessors().items()], reverse=rev)
#            successors_2 = sorted(successors_2, reverse=rev)
            utils, successors = zip(*successors)
#            print "utils=", utils
#            print "successors==?==successors_3", successors_3==successors
#            print successors[0], successors_3[0]
            if PRINT_TRACE_TYPE2: print "sorted"
        trim_limit = int(ceil(len(successors)*self.w))
#        trim_limit = max(trim_limit, 10)
        successors = successors[:trim_limit]
        utils = utils[:trim_limit]
#        print "SUC :",successors
        
        if values_only == True:
            successors = zip(*successors)[1]
        if PRINT_TRACE_TYPE0 and depth == 0 and self.max_depth == 1: print "player=%s, len successors=%s" % (q_state.curPlayer, len(successors)) 
        return successors, utils
    
    
    def _getSuccFn(self, state):
        if state.getCurrentPlayer() == self.player:
            return self._maxValue
        else:
            return self._minValue

    def get_successors_by_relevance(self, q_state, depth, next_rowW, next_rowB, rev=False, values_only=True):

        successors_dic = {}
        if not self.sorting_on or depth == self.max_depth - 1:
            successors = q_state.getSuccessors().items()              
        else:            
            successors = q_state.getSuccessors().items()  
            for action, q_state in successors:
                if action.type == 'M' or action.type == 'P':
                    successors_dic[(action, q_state)] = 0 
                else: 
                    if action.type == 'WH':
                        dis = map(lambda loc: (loc[0]-action.info[0],loc[1]-action.info[1]), q_state.locationsW + q_state.locationsB)
                    elif action.type == 'WV':
                        dis = map(lambda loc: (loc[0]-action.info[1],loc[1]-action.info[0]), q_state.locationsW + q_state.locationsB)
                    else:
                        raise Exception('bad move')
            
                    min_rate = INFINITY
                    for d in dis:
                        rate = 1
                        if (d[0] == 0 or  d[0] == 1):
                            rate += 0
                        elif (d[0] < 0):
                            rate += abs(d[0])
                        elif (d[0] > 1):
                            rate += abs(d[0]) - 1
                            
                        if (d[1] == 0 or  d[1] == 1):
                            rate += 0
                        elif (d[1] < 0):
                            rate += abs(d[1])
                        elif (d[1] > 1):
                            rate += abs(d[1]) - 1
                        
                        if rate < min_rate:
                            min_rate = rate
                            
                        successors_dic[(action, q_state)] = min_rate
                
            successors = list(zip(*sorted(successors_dic.items(), key=lambda suc: suc[1]))[0])
                
        trim_limit = int(ceil(len(successors)*self.w))
#        trim_limit = max(trim_limit, 10)
        successors = successors[:trim_limit]

        if values_only == True:
            successors = zip(*successors)[1]
        utils = [None]* len(successors)
        return successors, utils
    
